package 二叉树剪枝;

import tree_node.TreeNode;

class Solution {
    //初次尝试
    public TreeNode pruneTree1(TreeNode root) {
        dfs(root);
        if(root.left == null && root.right == null && root.val == 0){
            return null;
        }
        return root;
    }

    public boolean dfs(TreeNode root){
        if(root == null){
            return false;
        }
        boolean left = dfs (root.left);
        boolean right = dfs (root.right);
        if(left == false){
            root.left = null;
        }
        if(right == false){
            root.right = null;
        }
        if(left != false || right != false){
            return true;
        }
        return root.val == 1 ? true : false;
    }

    //进行优化
    public TreeNode pruneTree(TreeNode root) {
        if(root == null){
            return root;
        }
        root.left = pruneTree(root.left);
        root.right = pruneTree(root.right);
        if(root.val == 0 && root.left == null && root.right == null){
            return null;
        }
        return root;
    }
}